#include <iostream>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <cstring>
#include <cassert>
#include <cstdlib>
#include <sys/epoll.h>
#include <cerrno>
using namespace std;

const int MAX_EVENT_NUMBER = 1024; // 最大事件数
const int TCP_BUFFER_SIZE = 512; // TCP缓冲区大小
const int UDP_BUFFER_SIZE = 1024; // UDP缓冲区大小

int setNonBlocking(int fd) {
    int oldOption = fcntl(fd, F_GETFL);
    int newOption = oldOption | O_NONBLOCK;
    fcntl(fd, F_SETFL, newOption);
    return oldOption;
}

void addFd(int epollFd, int fd) {
    epoll_event event;
    event.data.fd = fd;
    event.events = EPOLLIN | EPOLLET; // 监听输入和边缘触发事件
    epoll_ctl(epollFd, EPOLL_CTL_ADD, fd, &event)
    setNonBlocking(fd); // 设置非阻塞模式
}

int main(int argc, char* argv[]) {
    if (argc <= 2) {
        cerr << "Usage: " << argv[0] << " <IP> <Port>" << endl;
        return 1;
    }
    const char* ip = argv[1];
    int port = atoi(argv[2]);

    int ret = 0;
    // 创建TCP套接字
    struct sockaddr_in address;
    memset(&address, 0, sizeof(address));
    address.sin_family = AF_INET;
    inet_pton(AF_INET, ip, &address.sin_addr);
    address.sin_port = htons(port);

    int tcpFd = socket(AF_INET, SOCK_STREAM, 0);
    assert(tcpFd >= 0);
    ret = bind(tcpFd, (struct sockaddr*)&address, sizeof(address));
    assert(ret != -1);
    ret = listen(tcpFd, 5);
    assert(ret != -1);

    // 创建UDP套接字
    memset(&address, 0, sizeof(address));
    address.sin_family = AF_INET;
    inet_pton(AF_INET, ip, &address.sin_addr);
    address.sin_port = htons(port);

    int udpFd = socket(AF_INET, SOCK_DGRAM, 0);
    assert(udpFd >= 0);
    ret = bind(udpFd, (struct sockaddr*)&address, sizeof(address));
    assert(ret != -1);
    // 在 UDP 中，服务器不需要建立连接，客户端可以直接向服务器发送数据报

    epoll_event events[MAX_EVENT_NUMBER];
    int epollFd = epoll_create(5);
    assert(epollFd != -1);
    addFd(epollFd, tcpFd); // 添加TCP套接字到epoll
    addFd(epollFd, udpFd); // 添加UDP套接字到epoll 
    // udp没有建立连接，所有的客户端都可以直接发送数据报，因此直接在这个fd上处理读写数据

    while (1)
    {
        int number = epoll_wait(epollFd, events, MAX_EVENT_NUMBER, -1);
        if (number < 0) {
            cerr << "epoll_wait error: " << strerror(errno) << endl;
            break;
        }

        for (int i = 0; i < eventCount; ++i) {
            int fd = events[i].data.fd;
            if(fd == tcpFd) {
                // 处理TCP连接
                struct sockaddr_in clientAddr;
                socklen_t clientAddrLen = sizeof(clientAddr);
                int connFd = accept(tcpFd, (struct sockaddr*)&clientAddr, &clientAddrLen);
                addFd(epollFd, connFd); // 添加新连接到epoll
            } 
            else if (fd == udpFd) {
                // 处理UDP数据报
                char buffer[UDP_BUFFER_SIZE];
                memset(buffer, 0, UDP_BUFFER_SIZE);
                struct sockaddr_in clientAddr;
                socklen_t clientAddrLen = sizeof(clientAddr);
                // UDP 两个专属 IO
                ret = recvfrom(udpFd, buffer, UDP_BUFFER_SIZE - 1, 0, 
                               (struct sockaddr*)&clientAddr, &clientAddrLen);
                // 收到什么就回什么
                if (ret > 0) {
                    sendto(udpFd, buffer, UDP_BUFFER_SIZE - 1, 0, 
                           (struct sockaddr*)&clientAddr, clientAddrLen);
                }
            } 
            else if (events[i].events & EPOLLIN) {
                // 处理TCP客户端发送的数据
                char buffer[TCP_BUFFER_SIZE];
                // ET模式必须循环读取，一次通知后读完
                while (1)
                {
                    memset(buffer, 0, TCP_BUFFER_SIZE);
                    ret = recv(fd, buffer, TCP_BUFFER_SIZE - 1, 0);
                    if (ret < 0) {
                        if(errno == EAGAIN || errno == EWOULDBLOCK) {
                            break;  // 非阻塞模式下没有数据可读
                        } 
                        close(fd);  // 这里是发生了其它的错误，
                        break;
                    } 
                    else if (ret == 0) {
                        close(fd);
                    }
                    else {
                        send (fd, buffer, ret, 0); // 回显数据
                    }
                }
            }
            else {
                cout << "Unknown event on fd: " << fd << endl;
            }
        }
    }

    close(tcpFd);
    close(udpFd);
    close(epollFd);
    return 0;
}